import struct
import torch
import torch.nn.functional as F
from config import ACTIVITY_LABELS
from data_processor import preprocess_data
from visualization import plot_sensor_data
from data_saver import save_data

def handle_client(client_socket, addr, model, device):
    """处理二进制格式的客户端连接"""
    print(f"接受来自 {addr} 的连接")

    # 使用二进制模式接收数据
    buffer = b""
    header_detected = False
    data_length = 0
    raw_data = []

    try:
        while True:
            # 接收二进制数据
            data = client_socket.recv(1024)
            if not data:
                break

            buffer += data

            # 检查是否有足够的数据来处理
            while len(buffer) > 0:
                # 如果还没有检测到头部，尝试查找头部标记
                if not header_detected:
                    # 尝试查找文本标记 "HAR_DATA\n"
                    try:
                        text_header_pos = buffer.find(b"HAR_DATA\n")
                        if text_header_pos >= 0:
                            # 找到文本头部，切换到文本处理模式
                            print("检测到文本数据格式")
                            return handle_text_client(client_socket, addr, model, device, buffer[text_header_pos:])
                    except:
                        pass

                    # 如果缓冲区至少有4个字节，检查二进制头部
                    if len(buffer) >= 4:
                        # 假设前4个字节是一个int32，表示数据包中的样本数
                        data_length = struct.unpack('<I', buffer[:4])[0]
                        print(f"检测到二进制数据格式，样本数: {data_length}")
                        buffer = buffer[4:]
                        header_detected = True
                    else:
                        # 数据不足，等待更多数据
                        break
                else:
                    # 已经检测到头部，现在处理数据
                    # 每个样本有6个浮点数 (3轴加速度 + 3轴陀螺仪)，每个浮点数4字节
                    bytes_per_sample = 6 * 4
                    total_bytes_needed = data_length * bytes_per_sample

                    if len(buffer) >= total_bytes_needed:
                        # 有足够的数据来处理所有样本
                        raw_data = []
                        for i in range(data_length):
                            offset = i * bytes_per_sample
                            sample_data = []
                            for j in range(6):
                                value = struct.unpack('<f', buffer[offset + j * 4:offset + (j + 1) * 4])[0]
                                sample_data.append(value)
                            raw_data.append(sample_data)

                        # 处理完数据后，清空缓冲区
                        buffer = buffer[total_bytes_needed:]
                        header_detected = False

                        print(f"接收完成，共 {len(raw_data)} 个样本")

                        # 处理数据并进行预测
                        if len(raw_data) > 0:
                            features = preprocess_data(raw_data)
                            features_tensor = torch.tensor([features], dtype=torch.float32).to(device)

                            with torch.no_grad():
                                outputs = model(features_tensor)
                                _, predicted = torch.max(outputs, 1)
                                predicted_label = predicted.item()

                                # 计算概率
                                probabilities = F.softmax(outputs, dim=1)[0]

                                print("\n===== 人类活动识别结果 =====")
                                print(f"预测活动: {ACTIVITY_LABELS[predicted_label]}")
                                print("\n各活动类别的概率:")
                                for i, prob in enumerate(probabilities):
                                    print(f"{ACTIVITY_LABELS[i]}: {prob.item():.4f} ({prob.item() * 100:.2f}%)")
                                print("=============================\n")
                                
                                # 可视化数据
                                plot_sensor_data(raw_data)
                                
                                # 保存数据
                                save_data(raw_data, predicted_label)
                    else:
                        # 数据不足，等待更多数据
                        break

    except Exception as e:
        print(f"处理客户端时出错: {e}")
    finally:
        client_socket.close()
        print(f"与 {addr} 的连接已关闭")


def handle_text_client(client_socket, addr, model, device, initial_buffer=b""):
    """处理文本格式的客户端连接"""
    print(f"切换到文本处理模式")

    buffer = initial_buffer.decode('utf-8', errors='ignore')
    collecting = False
    raw_data = []

    try:
        while True:
            try:
                data = client_socket.recv(1024).decode('utf-8', errors='ignore')
                if not data:
                    break

                buffer += data
            except UnicodeDecodeError:
                print("警告: 接收到无法解码的数据，尝试以二进制模式处理")
                return handle_client(client_socket, addr, model, device)

            lines = buffer.split('\n')
            buffer = lines.pop()  # 保留最后一个不完整的行

            for line in lines:
                if line == "HAR_DATA":
                    collecting = True
                    raw_data = []
                    print("开始接收新数据...")
                elif line == "END" and collecting:
                    collecting = False
                    print(f"接收完成，共 {len(raw_data)} 个样本")

                    # 处理数据并进行预测
                    if len(raw_data) > 0:
                        features = preprocess_data(raw_data)
                        features_tensor = torch.tensor([features], dtype=torch.float32).to(device)

                        with torch.no_grad():
                            outputs = model(features_tensor)
                            _, predicted = torch.max(outputs, 1)
                            predicted_label = predicted.item()

                            # 计算概率
                            probabilities = F.softmax(outputs, dim=1)[0]

                            print("\n===== 人类活动识别结果 =====")
                            print(f"预测活动: {ACTIVITY_LABELS[predicted_label]}")
                            print("\n各活动类别的概率:")
                            for i, prob in enumerate(probabilities):
                                print(f"{ACTIVITY_LABELS[i]}: {prob.item():.4f} ({prob.item() * 100:.2f}%)")
                            print("=============================\n")

                            # 可视化数据
                            plot_sensor_data(raw_data)

                            # 保存数据
                            save_data(raw_data, predicted_label)

                elif collecting:
                    try:
                        # 解析数据行 (accel_x,accel_y,accel_z,gyro_x,gyro_y,gyro_z)
                        values = [float(x) for x in line.split(',')]
                        if len(values) == 6:
                            raw_data.append(values)
                    except ValueError:
                        print(f"无法解析数据行: {line}")

    except Exception as e:
        print(f"处理文本客户端时出错: {e}")
    finally:
        client_socket.close()
        print(f"与 {addr} 的连接已关闭")